# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question3/electrostatic_solver.py
import numpy
import numpy.typing


def solve_dirichlet(
        potential: numpy.typing.NDArray[numpy.object_]) \
        -> numpy.typing.NDArray[numpy.float_]:
    """
    求解给定了狄拉克边界条件的二维无源静电势场。
    :param potential: 用以给定求解区域和边界条件。
                      在两个方向上每一单元代表的长度被视作相等的。
                      如果某点的电势是已知的，那么应该在此单元中填入具体的值，即边界条件；
                      如果是需要求解的，那么应该在此单元中填入 `None` 。
    :return: 求解结果。
    """
    row_count, column_count = potential.shape
    total_count = row_count * column_count
    flattened = potential.reshape([total_count])

    # TODO: 使用稀疏矩阵或其他方式进行优化
    a = numpy.zeros([total_count, total_count])
    b = numpy.zeros([total_count])
    for k in range(total_count):
        potential = flattened[k]
        if potential is not None:
            a[k, k] = 1
            b[k] = potential
        else:
            a[k, k] = 4
            a[k, k - 1] = -1
            a[k, k + 1] = -1
            a[k, k - column_count] = -1
            a[k, k + column_count] = -1

    flattened = numpy.linalg.solve(a, b)
    return flattened.reshape([row_count, column_count])
